from langchain_community.chat_message_histories import RedisChatMessageHistory, ChatMessageHistory

from langChain.config import model
from langchain_core.chat_history import BaseChatMessageHistory, InMemoryChatMessageHistory
from langchain_core.messages import HumanMessage, AIMessage, SystemMessage, trim_messages
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables import RunnableWithMessageHistory, RunnablePassthrough, ConfigurableFieldSpec

prompt = ChatPromptTemplate.from_messages([
    ("system", "你是一个乐于助人的助手，用{language}尽你所能回答所有问题。"),
    MessagesPlaceholder(variable_name="my_msg"),
    MessagesPlaceholder(variable_name="chat_history"),  # 历史消息key
])
# 3. 构建处理链
chain = prompt | model

# # 存储消息的历史记录
# store = {}
#
#
# def get_session_history(session_id: str) -> BaseChatMessageHistory:
#     if session_id not in store:
#         store[session_id] = InMemoryChatMessageHistory()
#     return store[session_id]
#
#
# # 模型的消息历史包装器
# with_message_history = RunnableWithMessageHistory(
#     chain,
#     get_session_history,
#     input_messages_key="my_msg",
#     history_messages_key="chat_history",
# )
#
# config = {"configurable": {"session_id": "abc2"}}
# response = with_message_history.invoke(
#     {'language': "中文", 'my_msg': [HumanMessage(content="你好啊，我是王东平")], },
#     config=config,
# )
# print(f'通过历史消息，第一次返回数据是：{response.content}')
#
# response = with_message_history.invoke(
#     {'language': "中文", 'my_msg': [HumanMessage(content="请问：我的名字是什么?")]},
#     config=config,
# )
#
# print(f'通过历史消息，第二次返回数据是：{response.content}')
#
# response = with_message_history.invoke(
#     {'language': "中文", 'my_msg': [HumanMessage(content="请问：我的名字是什么?")]},
#     config={"configurable": {"session_id": "abc3"}},
# )
#
# print(f'通过历史消息，引用不同的 session_id，返回数据是：{response.content}')
#
# print('=========================================================')

# # 扩展默认配置，同时配置user_id 和 session_id，甚至更多参数
# store_2 = {}
#
#
# def get_session_history(user_id: str, session_id: str) -> BaseChatMessageHistory:
#     if (user_id, session_id) not in store_2:
#         store_2[(user_id, session_id)] = InMemoryChatMessageHistory()
#     return store_2[(user_id, session_id)]
#
# # history_factory_config ：用于自定义会话标识的生成规则和扩展历史存储的查询维度。其核心作用是通过定义更灵活的会话标识方式，解决默认仅支持 session_id 的局限性，从而适应复杂业务场景的需求。
# # 默认情况下，RunnableWithMessageHistory 仅支持通过 session_id（字符串）标识会话。
# # 通过 history_factory_config，可定义多字段组合标识会话（如 user_id + conversation_id），使会话管理更贴合业务逻辑。
# with_message_history = RunnableWithMessageHistory(
#     chain,
#     get_session_history,
#     input_messages_key="my_msg",
#     history_messages_key="chat_history",
#     history_factory_config=[
#         ConfigurableFieldSpec(
#             id="user_id",
#             annotation=str,
#             name="user_id",
#             description="用户id",
#             default="",
#             is_shared=True,
#         ),
#         ConfigurableFieldSpec(
#             id="session_id",
#             annotation=str,
#             name="session_id",
#             description="用户会话id",
#             default="",
#             is_shared=True,
#         ),
#     ]
# )
#
# config = {"configurable": {"user_id": "user1", "session_id": "abc2"}}
# response = with_message_history.invoke(
#     {'language': "中文", 'my_msg': [HumanMessage(content="你好啊，我是王东平")], },
#     config=config,
# )
# print(f'通过历史消息，第一次返回数据是：{response.content}')
#
# response = with_message_history.invoke(
#     {'language': "中文", 'my_msg': [HumanMessage(content="请问：我的名字是什么?")]},
#     config=config,
# )
#
# print(f'通过历史消息，第二次返回数据是：{response.content}')
#
# response = with_message_history.invoke(
#     {'language': "中文", 'my_msg': [HumanMessage(content="请问：我的名字是什么?")]},
#     config={"configurable": {"user_id": "user1", "session_id": "abc3"}},
# )
#
# print(f'通过历史消息，引用不同的 session_id，返回数据是：{response.content}')
# print('=========================================================')

# # 把历史消息存入redis
# REDIS_URL = "redis://localhost:6379/0"
#
#
# def get_session_history_by_redis(session_id: str) -> RedisChatMessageHistory:
#     return RedisChatMessageHistory(session_id, url=REDIS_URL)
#
#
# with_message_history_by_redis = RunnableWithMessageHistory(
#     chain,
#     get_session_history=get_session_history_by_redis,
#     input_messages_key="my_msg",
#     history_messages_key="chat_history",
# )
#
# config = {"configurable": {"session_id": "abc2"}}
# response = with_message_history_by_redis.invoke(
#     {'language': "中文", 'my_msg': [HumanMessage(content="你好啊，我是王东平")], },
#     config=config,
# )
# print(f'通过历史消息，第一次返回数据是：{response.content}')
#
# response = with_message_history_by_redis.invoke(
#     {'language': "中文", 'my_msg': [HumanMessage(content="请问：我的名字是什么?")]},
#     config=config,
# )
#
# print(f'通过历史消息，第二次返回数据是：{response.content}')
#
# response = with_message_history_by_redis.invoke(
#     {'language': "中文", 'my_msg': [HumanMessage(content="请问：我的名字是什么?")]},
#     config={"configurable": {"session_id": "abc3"}},
# )
#
# print(f'通过历史消息，引用不同的 session_id，返回数据是：{response.content}')
# print('=========================================================')

# 消息裁剪，即保留最近的 n条消息
temp_chat_history = ChatMessageHistory()
temp_chat_history.add_user_message("我叫Jack，你好")
temp_chat_history.add_ai_message("你好")
temp_chat_history.add_user_message("我今天心情挺开心")
temp_chat_history.add_ai_message("你今天心情怎么样")
temp_chat_history.add_user_message("我下午在打篮球")
temp_chat_history.add_ai_message("你下午在做什么")

prompt = ChatPromptTemplate.from_messages([
    ("system", "你是一个乐于助人的助手。尽力回答所有问题。提供的聊天历史包括与您交谈的用户的事实。",),
    MessagesPlaceholder(variable_name="chat_history"),
    ("human", "{input}"),
])
chain_with_message_history = RunnableWithMessageHistory(
    prompt | model,
    lambda session_id: temp_chat_history,
    input_messages_key="input",
    history_messages_key="chat_history",
)


# response = chain_with_message_history.invoke(
#     {"input": [HumanMessage(content="我今天心情如何?")]},
#     {"configurable": {"session_id": "unused"}},
# )
# print(response)
# print(temp_chat_history.messages)


# 裁剪消息
def trim_messages(chain_input):
    stored_messages = temp_chat_history.messages
    if len(stored_messages) <= 2:
        return False
    temp_chat_history.clear()
    for message in stored_messages[-2:]:
        temp_chat_history.add_message(message)
    return True


chain_with_trimming = RunnablePassthrough.assign(messages_trimmed=trim_messages) | chain_with_message_history

response_1 = chain_with_trimming.invoke(
    {"input": [HumanMessage(content="我下午在干嘛?")]},
    {"configurable": {"session_id": "unused"}},
)
print('AI的回答 :', response_1.content, '\n 保留的历史消息:', temp_chat_history, '\n\n')

response_2 = chain_with_trimming.invoke(
    {"input": [HumanMessage(content="我的心情怎么样?")]},
    {"configurable": {"session_id": "unused"}},
)
print('AI的回答 :', response_2.content, '\n 保留的历史消息:', temp_chat_history, '\n\n')

# # 总结记忆
# def summarize_messages(chain_input):
#     stored_messages = temp_chat_history.messages
#     if len(stored_messages) == 0:
#         return False
#     summarization_prompt = ChatPromptTemplate.from_messages([
#         MessagesPlaceholder(variable_name="chat_history"),
#         ("user", "将上述聊天消息浓缩成一条摘要消息。尽可能包含多个具体细节。",),
#     ])
#     summarization_chain = summarization_prompt | model
#     summary_message = summarization_chain.invoke({"chat_history": stored_messages})
#     temp_chat_history.clear()
#     temp_chat_history.add_message(summary_message)
#     return True
#
#
# chain_with_summarization = RunnablePassthrough.assign(
#     messages_summarized=summarize_messages) | chain_with_message_history
#
# response = chain_with_summarization.invoke(
#     {"input": "名字，下午在干嘛，心情"},
#     {"configurable": {"session_id": "unused"}},
# )
# print(response.content, '\n')
#
# print('生成的摘要信息 ： ', temp_chat_history.messages)
